import torch
import matplotlib.pyplot as plt

"""
在momentum动量法中，动量项被引入来模拟物理学中的动量概念，即物体在它运动方向上保持运动的趋势。
mt = βmt-1 + (1-β)gt
当前时刻的动量 = β*上一时刻累积的动量 + (1-β)*当前时刻的梯度
"""

# 1. 定义损失函数
def loss_fn(w1, w2):
    return w1 ** 2 + 2 * w2 ** 2

# 2. 超参数
lr = 0.05
Epochs = 20
beta = 0.9
w1 = -1
w2 = 1
v1 = 0
v2 = 0


''' 绘制等高线图 '''
x1 = torch.linspace(-1, 1, 100)
x2 = torch.linspace(-1, 1, 100)
xx1, xx2 = torch.meshgrid(x1, x2, indexing="ij")
loss = loss_fn(xx1, xx2)
fig = plt.figure("GD with Momentum")
ax = fig.add_subplot()
ax.contour(xx1, xx2, loss)
# 定义一个列表，用于存储梯度下降的路径点
points = []

# 3. 循环训练
for epoch in range(Epochs):
    points.append([w1, w2])     # 保存参数点
    loss = loss_fn(w1, w2)
    print(loss)
    # ==== GD with Momentum算法 ====
    # 计算梯度
    g1 = 2 * w1
    g2 = 4 * w2
    # 计算动量
    v1 = beta * v1 + (1 - beta) * g1
    v2 = beta * v2 + (1 - beta) * g2
    # 更新参数
    w1 -= lr * v1
    w2 -= lr * v2

points = torch.tensor(points)
ax.plot(points[:,0], points[:,1], "ko-")
plt.show()
